package fons.navigator.views.island.maps;

import java.awt.Color;
import java.util.Collection;
import java.util.Collections;
import java.util.List;
import java.util.Map;
import java.util.Set;

import com.google.common.collect.Lists;
import com.google.common.collect.Maps;
import com.google.common.collect.Sets;

import edu.umd.cs.piccolo.activities.PActivity;
import edu.umd.cs.piccolo.event.PBasicInputEventHandler;
import edu.umd.cs.piccolo.event.PInputEvent;
import fons.model.ontology.ValueSpace;
import fons.navigator.util.piccolo.PButton;
import fons.navigator.views.island.ColorField;
import fons.navigator.views.island.Island;
import fons.navigator.views.island.MapColorManager;
import fons.navigator.views.island.MapGrid;
import fons.navigator.views.island.MapModel;
import fons.navigator.views.island.VisualIsland;

public class NeuralMapModel extends MapModel {
	private final int SIZE;
	private final int DIMENSIONALITY;
	private final int MAX_STEP;
	private final double LEARNING_RATE;

	private final VisualIsland visualIsland;
	private final List<Neuron> neurons;
	private final Map<Neuron, Island> bmus = Maps.newHashMap();

	private int step = 1;
	private double offsetX;
	private double offsetY;

	private final double lambda;
	private ColorField colorField;

	private enum ViewType { CLUSTER, UMATRIX, MULTIPLY };
	private ViewType viewType = ViewType.CLUSTER;

	private class NeuralColorField implements ColorField {

		public NeuralColorField() {
			for (Neuron neuron: neurons) {
				Set<Neuron> neighbours = getNeighbours(neuron, 2);
				neighbours.remove(neuron);

				double scale = 0;
				for (Neuron neighbour: neighbours) {
					scale += neuron.weightDistance(neighbour);
				}
				scale /= getVisualIsland().getDimensionality();
				neuron.setScale(scale);
			}
		}

		@Override
		public Color at(double x, double y, final MapGrid.MapQuad q) {
			if (x > 1.5 || x < -0.5 || y > 1.5 || y < - 0.5) return Color.BLACK;

			boolean dampening = false;
			if (x > 1) {
				x = x - (int) x;
				dampening = true;
			}
			if (x < 0) {
				x = -x;
				x = 1 - (x - (int) x);
				dampening = true;
			}
			if (y > 1) {
				y = y - (int) y;
				dampening = true;
			}
			if (y < 0) {
				y = -y;
				y = 1 - (y - (int) y);
				dampening = true;
			}

			int nx = (int) (x * SIZE + offsetX * SIZE) % SIZE;
			int ny = (int) (y * SIZE + offsetY * SIZE) % SIZE;

			final Neuron n = getNeuron(nx, ny);
			if (bmus.containsKey(n)) {
				Island island = bmus.get(n);
				final double quadWidth = visualIsland.getMapGrid().getQuadWidth();
				final double quadHeight = visualIsland.getMapGrid().getQuadHeight();
				final double scale = visualIsland.getValueScale();
				island.placeAt(x + quadWidth*scale/2,
						y + quadHeight*scale/2);
			}

			float value = 1 - (float) (n.getScale() * getVisualIsland().getColorIntensity());
			if (viewType == ViewType.CLUSTER)
				value *= getVisualIsland().getColorIntensity();
			if (value > 1) value = 1;
			else if (value < 0) value = 0;


			if (viewType == ViewType.UMATRIX) {
				return dampening ? 
						new Color(value/2, value/2, value/2)
				: new Color(value, value, value);
			}

			float green = 0;
			float red = 0; 
			float blue = 0;  

			final Collection<ValueSpace> valueSpaces = getVisualIsland().getValueSpaces();
			final MapColorManager colorManager = getVisualIsland().getColorManager();
			int i = 0;
			for (ValueSpace vsp: valueSpaces) {
				final float weight = (float) n.getWeight(i)*n.getWeight(i);

				red += colorManager.redOf(vsp, weight) * (float) (viewType == ViewType.MULTIPLY ? 
						getVisualIsland().getColorIntensity() : 1);
				green += colorManager.greenOf(vsp, weight) * (float) (viewType == ViewType.MULTIPLY ? 
						getVisualIsland().getColorIntensity() : 1);
				blue += colorManager.blueOf(vsp, weight) * (float) (viewType == ViewType.MULTIPLY ? 
						getVisualIsland().getColorIntensity() : 1);

				if (red > 1) red = 1; if (green > 1) green = 1; if (blue > 1) blue = 1;
				++i;
			}

			if (viewType == ViewType.MULTIPLY) {
				red *= value;
				green *= value;
				blue *= value;
			}

			return dampening ?
					new Color(red/2, green/2, blue/2) :
						new Color(red, green, blue);
		}

	}

	public NeuralMapModel (int size, int dimensionality, int steps, double learningRate, VisualIsland island) {
		super(island);

		this.offsetX = size/2.0;
		this.offsetY = size/2.0;

		island.getCamera().addInputEventListener(new PBasicInputEventHandler() {
			@Override
			public void mouseDragged(PInputEvent event) {
				if (event.isRightMouseButton()) {
					offsetX -= event.getCanvasDelta().width * getVisualIsland().getValueScale();
					offsetY -= event.getCanvasDelta().height * getVisualIsland().getValueScale();
					getVisualIsland().getMapGrid().updateColors();
				}
			}
		});

		for (Island i: island.getIslands()) {
			i.setOffset(Math.random()/island.getValueScale(), Math.random()/island.getValueScale());
		}

		this.SIZE = size;
		this.DIMENSIONALITY = dimensionality;
		this.MAX_STEP = steps;
		this.LEARNING_RATE = learningRate;
		this.neurons = Lists.newArrayListWithCapacity(SIZE*SIZE);
		this.visualIsland = island;

		this.lambda = MAX_STEP / Math.log(SIZE);

		for (int x = 0; x < SIZE; ++x) {
			for (int y = 0; y < SIZE; ++y) {
				this.neurons.add(new Neuron(x, y, SIZE, DIMENSIONALITY, island));
			}
		}
	}

	public Neuron getNeuron(int x, int y) {
		Neuron n = this.neurons.get(x*SIZE + y);
		assert(n.getX() == x && n.getY() == y);
		return n;
	}

	public List<Neuron> getNeurons() {
		return Collections.unmodifiableList(this.neurons);
	}

	public Neuron getBMU(Island island) {
		Neuron bmu = null;
		double distance = Double.NaN;

		for (Neuron neuron: neurons) {
			if (bmu == null) {
				distance = neuron.distance(island);
				bmu = neuron;
			} else {
				final double dist = neuron.distance(island);
				if (dist < distance) {
					bmu = neuron;
					distance = dist;
				}
			}
		}

		return bmu;
	}

	private Set<Neuron> getNeighbours(Neuron neuron, double neighbourhoodRadius) {
		Set<Neuron> neighbours = Sets.newHashSet();
		int xSpan = (int) neighbourhoodRadius;
		int ySpan = xSpan;
		
		/*for (Neuron n: neurons) {
			if (n.distance(neuron) < neighbourhoodRadius) {
				neighbours.add(n);
			}
		}*/
		
		for (int x = -xSpan; x < xSpan; ++x) {
			for (int y = -xSpan; y < ySpan; ++y) {
				
				int nx = neuron.getX() + x;
				if (nx >= SIZE) nx %= SIZE;
				else if (nx < 0) nx += SIZE;
				
				int ny = neuron.getY() + y;
				if (ny >= SIZE) ny %= SIZE;
				else if (ny < 0) ny += SIZE;
				
				neighbours.add(getNeuron(nx, ny));
			}
		}
		
		return neighbours;
	}
	
	private void step () {
		final double neighbourhoodRadius = SIZE/2 * Math.exp(-step / lambda);
		final double learningRate = LEARNING_RATE * Math.exp(-step / lambda);
				
		if (neighbourhoodRadius < 1) {
			return;
		}
		
		bmus.clear();
		for (Island island: visualIsland.getIslands()) {
			Neuron bmu = getBMU(island);
			bmus.put(bmu, island);
			learn(island, bmu, getNeighbours(bmu, neighbourhoodRadius), neighbourhoodRadius, learningRate);
		}

		step++;
	}

	private void learn(Island island, Neuron bmu, Set<Neuron> neighbours,
			double neighbourhoodRadius, double learningRate) {
		for (Neuron n: neighbours) {
			final double distance = bmu.distance(n);
			n.learn(island, Math.exp(-(distance*distance)/(2*neighbourhoodRadius*neighbourhoodRadius)) * learningRate);
		}	
	}

	@Override
	public void step(double time) {
		step();
	}

	@Override
	public ColorField getColorField() {
		if (colorField == null) updateColorField();
		return colorField;
	}

	@Override
	public void updateColorField() {
		this.colorField = new NeuralColorField();
		getVisualIsland().getMapGrid().updateColors();
	}

	@Override
	public void createButtons() {					 
		final PButton viewTypeButton = new PButton("      U-Matrix      ");
		getButtonsPanel().addButton(viewTypeButton, new PBasicInputEventHandler() {
			@Override
			public void mouseClicked(PInputEvent event) {
				switch (viewType) {
				case CLUSTER:
					viewTypeButton.setText(" U-Matrix + Cluster ");
					viewType = ViewType.UMATRIX;
					break;
				case UMATRIX:
					viewTypeButton.setText("       Cluster      ");
					viewType = ViewType.MULTIPLY;
					break;
				default:
					viewTypeButton.setText("      U-Matrix      ");
					viewType = ViewType.CLUSTER;
				}
				getVisualIsland().getMapGrid().updateColors();
			}
		});		

		getButtonsPanel().addButton(new PButton(" Step "), new PBasicInputEventHandler() {
			@Override
			public void mouseClicked(PInputEvent event) {
				getVisualIsland().addActivity(new PActivity(4000) {
					private long lastTime = 0;
					private long counter = 0;
					private long colorCounter = 0;
					@Override
					protected void activityStep(long elapsedTime) {
						final long dt = elapsedTime - lastTime;

						counter += dt;
						while (counter > 10) {
							counter -= 10;
							step();
						}

						colorCounter += dt;
						if (colorCounter > 250) {
							colorCounter = 0;
							updateColorField();
							getVisualIsland().getMapGrid().updateColors();
						}

						lastTime = elapsedTime;
					}
				});
			}
		});
	}

}
